package cn.jbolt.ai.core.chain.node.base;

import cn.hutool.core.util.StrUtil;
import cn.jbolt.ai.app.service.AiAppMessageService;
import cn.jbolt.ai.app.service.AiAppService;
import cn.jbolt.ai.app.service.AiAppSessionService;
import cn.jbolt.ai.core.chain.ChainContext;
import cn.jbolt.util.JSONUtil;
import cn.jbolt.util.ObjectUtil;
import com.fasterxml.jackson.databind.JsonNode;
import com.jboltai.capability.message.AIMessage;
import com.jboltai.event.Event;
import com.jboltai.event.ai.aichat.AIChatEvent;
import com.jboltai.event.ai.aitext.AITextEvent;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

@Component
public class BaseNode {


    @Autowired
    private AiAppService aiAppService;
    @Autowired
    private AiAppSessionService aiAppSessionService;
    @Autowired
    private AiAppMessageService aiAppMessageService;

    /**
     * 将值转换为字符串,普通类型直接tostring，复合类型转json字符串
     * @param value
     * @return
     */
    protected String valueToString(Object value) {
        if(ObjectUtil.isCommonTypeValue(value)){
            return Objects.toString(value);
        }
        return JSONUtil.toJson(value);
    }

    /**
     * 生成提示词
     *
     * @param config  配置
     * @param context 上下文
     * @return 处理后的提示词
     */
    public String generatePrompt(LlmConfig config, ChainContext context) {

        // 输入字段来自其他输出，则根据key读取对应值
        if (Objects.equals(config.getInputField(), AppConfigConstant.INPUT_FROM_CUSTOM) == false) {
            Object value = context.getAttr(config.getInputField());
            return valueToString(value);
        }

        // 否则就是自定义输入，就要根据模板渲染了
        // 模板为空，则直接返回原始的提示词
        if (StrUtil.isBlank(config.getPromptTpl())) {
            throw new RuntimeException("未设置提示词模版内容");
        }

        //模板不为空，则根据模板渲染
        String template = config.getPromptTpl();
        if (context.getArgs() == null || context.getArgs().isEmpty()) {
            return template;
        }


        return replaceArgs(template, context.getArgs());
    }

    private static final Pattern TEMPLATE_PATTERN = Pattern.compile("\\$\\{\\s*([^{}]+?)\\s*\\}");

    protected String replaceArgs(String template, Map<String, Object> args) {
        Matcher matcher = TEMPLATE_PATTERN.matcher(template);
        while (matcher.find()) {
            String key = matcher.group(1);
            Object value = args.get(key);
            if (value != null) {
                //直接取值
                // 处理可能存在的空格情况，如 ${ key } 或 ${key}
                template = template.replaceAll("\\$\\{\\s*" + key + "\\s*\\}", valueToString(value));
            } else if (key.contains(".")) {
                //key是类似 a.b[0] 这种格式，需要先根据a 从args中取出对象，再从对象身上根据 b[0] 取值
                String[] path = key.split("\\.", 2);
                value = args.get(path[0]);
                if (value == null) {
                    continue;
                }
                if (value instanceof JsonNode) {
                    template =  template.replaceAll("\\$\\{\\s*" + key + "\\s*\\}", JSONUtil.readValueAsText((JsonNode) value, path[1]));
                } else {
                    template =  template.replaceAll("\\$\\{\\s*" + key + "\\s*\\}", JSONUtil.readValueAsText(JSONUtil.toJsonNode(JSONUtil.toJson(value)), path[1]));
                }
            }
        }


        return template;
    }


    /**
     * 配置AI 相关参数
     *
     * @param config
     * @param context
     * @param event
     */
    public void configLlm(LlmConfig config, ChainContext context, Event event) {
        Long appId = context.getAiApp().getId();
        Long sessionId = context.getAiSession().getId();


        if (event instanceof AITextEvent) {
            AITextEvent textEvent = (AITextEvent) event;
            textEvent.prompt(generatePrompt(config, context));
            if (config.getModelType() != null) {
                textEvent.setModel(config.getModelType());
            }
            textEvent.getAiChatEvent().setStream(Objects.equals(config.isStream(), true));

            if (StrUtil.isNotBlank(config.getDefaultModel())) {
                textEvent.setModelName(config.getDefaultModel());
            }
            if (config.getMaxTokens() > 0) {
                textEvent.setMaxTokens(config.getMaxTokens());
            }
            if (StrUtil.isNotBlank(config.getSystemRole())) {
                textEvent.systemMessage(config.getSystemRole());
            }
            if (Objects.equals(config.isUseContext(), true)) {
                List<AIMessage> messageList = aiAppService.getSessionContextMessages(sessionId, context.getRequestMsg().getId(), null);
                if (messageList != null && messageList.size() > 0) {
                    textEvent.contextMessages(messageList);
                }
            }
            if (config.getTemperature() != null) {
                textEvent.setTemperature(config.getTemperature());
            }
            if (config.getTopP() != null && config.getTopP() > 0) {
                textEvent.setTopP(config.getTopP());
            }
            if (config.getTopK() != null && config.getTopK() > 0) {
                textEvent.setTopK(config.getTopK());
            }

        } else if (event instanceof AIChatEvent) {
            AIChatEvent aiChatEvent = (AIChatEvent) event;
            aiChatEvent.prompt(generatePrompt(config, context));
            aiChatEvent.setStream(Objects.equals(config.isStream(), true));
            if (config.getModelType() != null) {
                aiChatEvent.setModel(config.getModelType());
            }
            if (StrUtil.isNotBlank(config.getDefaultModel())) {
                aiChatEvent.setModelName(config.getDefaultModel());
            }
            if (config.getMaxTokens() > 0) {
                aiChatEvent.setMaxTokens(config.getMaxTokens());
            }
            if (StrUtil.isNotBlank(config.getSystemRole())) {
                aiChatEvent.systemMessage(config.getSystemRole());
            }
            if (Objects.equals(config.isUseContext(), true)) {
                List<AIMessage> messageList = aiAppService.getSessionContextMessages(sessionId, context.getRequestMsg().getId(), null);
                if (messageList != null && messageList.size() > 0) {
                    aiChatEvent.contextMessages(messageList);
                }
            }
            if (config.getTemperature() != null && config.getTemperature() > 0) {
                aiChatEvent.setTemperature(config.getTemperature());
            }
            if (config.getTopP() != null && config.getTopP() > 0) {
                aiChatEvent.setTopP(config.getTopP());
            }
            if (config.getTopK() != null && config.getTopK() > 0) {
                aiChatEvent.setTopK(config.getTopK());
            }
        }


    }


}
